package cn.bugstack.openai.executor.model.baichuan;

import cn.bugstack.openai.executor.Executor;
import cn.bugstack.openai.executor.model.baichuan.config.BaiChuanModelConfig;
import cn.bugstack.openai.executor.model.baichuan.valobj.BaiChuanModelCompletionRequest;
import cn.bugstack.openai.executor.model.baichuan.valobj.BaiChuanModelCompletionResponse;
import cn.bugstack.openai.executor.model.baichuan.valobj.FinishReason;
import cn.bugstack.openai.executor.parameter.*;
import cn.bugstack.openai.executor.result.ResultHandler;
import cn.bugstack.openai.session.Configuration;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.apache.commons.lang3.ObjectUtils;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.List;

/**
 * @author llq
 * @Description 百川 执行器
 */

@Slf4j
public class BaiChuanModelExecutor implements Executor, ResultHandler, ParameterHandler<BaiChuanModelCompletionRequest> {

    private EventSource.Factory factory;
    private BaiChuanModelConfig baiChuanModelConfig;

    public BaiChuanModelExecutor(Configuration configuration) {
        this.factory = configuration.createRequestFactory();
        this.baiChuanModelConfig = configuration.getBaiChuanModelConfig();
    }

    @Override
    public EventSource completions(CompletionRequest completionRequest, EventSourceListener eventSourceListener) throws Exception {
        // 1. 转换参数信息
        BaiChuanModelCompletionRequest baiChuanModelCompletionRequest = getParameterObject(completionRequest);
        // 2. 构建请求信息
        Request request = new Request.Builder()
                .header("Content-Type", "application/json")
                .header("Authorization", "Bearer " + baiChuanModelConfig.getApiKey())
                .url(baiChuanModelConfig.getApiHost().concat(baiChuanModelConfig.getV1_completions()))
                .post(RequestBody.create(MediaType.parse(Configuration.APPLICATION_JSON), JSON.toJSONString(baiChuanModelCompletionRequest)))
                .build();

        // 3. 返回事件结果
        return factory.newEventSource(request, eventSourceListener(eventSourceListener));
    }

    @Override
    public EventSource completions(String apiHostByUser, String apiKeyByUser, CompletionRequest completionRequest, EventSourceListener eventSourceListener) throws Exception {
        return null;
    }

    @Override
    public BaiChuanModelCompletionRequest getParameterObject(CompletionRequest completionRequest) {
        BaiChuanModelCompletionRequest baiChuanModelCompletionRequest = new BaiChuanModelCompletionRequest();
        baiChuanModelCompletionRequest.setModel(completionRequest.getModel());
        List<Message> messages = completionRequest.getMessages();
        List<cn.bugstack.openai.executor.model.baichuan.valobj.Message> baiChuanMessages = new ArrayList<>();
        for (Message message : messages) {
            baiChuanMessages.add(cn.bugstack.openai.executor.model.baichuan.valobj.Message.builder()
                    .role(message.getRole())
                    .content(message.getContent())
                    .build());
        }
        baiChuanModelCompletionRequest.setMessages(baiChuanMessages);


        return baiChuanModelCompletionRequest;
    }

    @Override
    public EventSourceListener eventSourceListener(EventSourceListener eventSourceListener) {

        return new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                BaiChuanModelCompletionResponse response = JSON.parseObject(data, BaiChuanModelCompletionResponse.class);
                // 正在生成
                List<BaiChuanModelCompletionResponse.Choice> choices = response.getChoices();
                if (ObjectUtils.isEmpty(choices.get(choices.size()-1).getFinish_reason())) {
                    CompletionResponse completionResponse = new CompletionResponse();

                    ArrayList<ChatChoice> chatChoices = new ArrayList<>();
                    ChatChoice chatChoice = new ChatChoice();
                    chatChoice.setDelta(Message.builder()
                            .role(CompletionRequest.Role.SYSTEM)
                            .name("")
                            .content(choices.get(choices.size()-1).getDelta().getContent())
                            .build());
                    chatChoices.add(chatChoice);
                    completionResponse.setChoices(chatChoices);

                    eventSourceListener.onEvent(eventSource, id, type, JSON.toJSONString(completionResponse));
                } else if (FinishReason.STOP.getCode().equals(choices.get(choices.size()-1).getFinish_reason())) { // 生成结束
                    ArrayList<ChatChoice> chatChoices = new ArrayList<>();
                    ChatChoice chatChoice = new ChatChoice();
                    chatChoice.setFinishReason("stop");
                    chatChoice.setDelta(Message.builder()
                            .role(CompletionRequest.Role.SYSTEM)
                            .name("")
                            .content(choices.get(choices.size()-1).getDelta().getContent())
                            .build());
                    chatChoices.add(chatChoice);
                    // 构建结果
                    CompletionResponse completionResponse = new CompletionResponse();
                    completionResponse.setChoices(chatChoices);
                    completionResponse.setCreated(System.currentTimeMillis());
                    // 返回数据
                    eventSourceListener.onEvent(eventSource, id, type, JSON.toJSONString(completionResponse));
                } else {
                    onClosed(eventSource);
                }


            }

            @Override
            public void onOpen(EventSource eventSource, Response response) {
                eventSourceListener.onOpen(eventSource, response);
            }



            @Override
            public void onClosed(EventSource eventSource) {
                eventSourceListener.onClosed(eventSource);
            }

            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                eventSourceListener.onFailure(eventSource, t, response);
            }
        };
    }
}
